import torch
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from torch.nn import Linear, Module

dataset = torchvision.datasets.CIFAR10(root="./dataset",
                                       train=False,
                                       transform=transforms.ToTensor(),
                                       download=True)
# drop_last：丢弃最后不满足batch_size=64的数据，不使用的话会造成最后一组数据和线性层匹配不上
dataloader = DataLoader(dataset, batch_size=64, shuffle=True, drop_last=True)

class MyLinear(Module):
    def __init__(self):
        super(MyLinear, self).__init__()
        self.linear = Linear(in_features=196608, out_features=10)

    def forward(self, x):
        return self.linear(x)

mylinear = MyLinear()

for data in dataloader:
    imgs, targets = data
    # 这样使用没有问题
    # imgs = imgs.reshape((1, 1, 1, -1))
    # imgs = imgs.reshape((1, -1))
    # 或是使用flatten()方法,将数据摊平成一行
    imgs = torch.flatten(imgs)
    print(imgs.shape)
    x = mylinear(imgs)
    print(x.shape)